package controller

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"github.com/gin-gonic/gin"
	"log"
)

// GenerateKeyPair generates an RSA key pair of the specified bit size
func GenerateKeyPair(bits int) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	priv, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, nil, err
	}
	return priv, &priv.PublicKey, nil
}

// ExportPublicKeyPEM exports an RSA public key to PEM format
func ExportPublicKeyPEM(pub *rsa.PublicKey) string {
	pubASN1, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return ""
	}

	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	})

	return string(pubPEM)
}

// ExportPrivateKeyPEM exports an RSA private key to PEM format
func ExportPrivateKeyPEM(priv *rsa.PrivateKey) string {
	privASN1 := x509.MarshalPKCS1PrivateKey(priv)

	privPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privASN1,
	})

	return string(privPEM)
}

// Encrypt1 encrypts plaintext with RSA public key
func Encrypt1(pub *rsa.PublicKey, plaintext []byte) ([]byte, error) {
	label := []byte("") // Optional label for OAEP padding, can be nil if not used
	ciphertext, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pub, plaintext, label)
	if err != nil {
		return nil, err
	}
	return ciphertext, nil
}

// Decrypt decrypts ciphertext with RSA private key
func Decrypt(priv *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	label := []byte("") // Must match the label used during encryption
	plaintext, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, priv, ciphertext, label)
	if err != nil {
		return nil, err
	}
	return plaintext, nil
}

// EncryptWithRSAPublicKey 使用RSA公钥加密数据
func EncryptWithRSAPublicKey(publicKeyPEM string, message []byte) (string, error) {
	// 解码PEM格式的公钥
	block, _ := pem.Decode([]byte(publicKeyPEM))
	if block == nil || block.Type != "PUBLIC KEY" {
		return "", errors.New("failed to decode PEM block containing public key")
	}

	// 解析公钥
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return "", err
	}

	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return "", errors.New("not a valid RSA public key")
	}

	// 使用OAEP填充和SHA-256哈希函数进行加密
	hash := sha256.New()
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, rsaPub, message, nil)
	if err != nil {
		return "", err
	}

	// 将密文编码为base64字符串
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func Test2(c *gin.Context) {
	// 示例公钥（PEM格式）
	publicKeyPEM := `-----BEGIN PUBLIC KEY-----  
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCt7bpQlP7sI0pqgrilzKgZx09s
VzzNXsEp0JOW1QUre3ifMt4R2SWdpYw9Lu0zoJUAIcO3jG43gfKzG2gP3HjmaFib
sgPjw2gxmSV0zY0MAN26gP+sRC5AEnlEz+goSS3JLz2TdMsTYiG0wo053O/a4IdI
nFDuE7JYnEx2po18QwIDAQAB
-----END PUBLIC KEY-----`
	data := gin.H{
		"id_card":   "11111",
		"card_type": 2,
		"name":      "333",
		//"birthday":   "444",
		//"sex":        "555",
		//"phone":      "666",
		//"address":    "777",
		//"drive_type": "888",
		//"portrait":   "999",
		//"inscode":    "10101010",
		//"reg_time":   "11111111",
		//"third_code": "12121212",
	}
	message, err := json.Marshal(data)
	if err != nil {
		return
	}
	//block, _ := pem.Decode([]byte(publicKeyPEM))
	//if block == nil || block.Type != "PUBLIC KEY" {
	//	log.Println("1111111111111111111")
	//
	//}
	//
	//// 解析公钥
	//pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	//if err != nil {
	//	log.Println("22222222222222222")
	//}
	//
	//rsaPub, _ := pub.(*rsa.PublicKey)
	////log.Println(ok)
	//encryptedBlocks, err := EncryptLongMessageRSA(rsaPub, message)
	//if err != nil {
	//	log.Println("Error encrypting long message:", err)
	//	return
	//}
	//log.Println(encryptedBlocks)
	//// 要加密的消息
	//message1 := []byte("Hello, RSA encryption! old be deng")
	//log.Println(message)
	//log.Println(message1)
	// 使用公钥加密消息
	encryptedMessage, err := EncryptWithRSAPublicKey(publicKeyPEM, message)
	if err != nil {
		log.Printf("Encryption error: %v\n", err)
		return
	}

	// 输出加密后的消息（base64编码）
	log.Printf("Encrypted message: %s\n", encryptedMessage)
}

const maxMessageSize = 128

func EncryptLongMessageRSA(publicKey *rsa.PublicKey, message []byte) ([]string, error) {
	var encryptedBlocks []string
	for len(message) > 0 {
		block := message[:maxMessageSize]
		if len(block) > len(message) {
			block = message
		}
		encryptedBlock, err := EncryptRSA(publicKey, block)
		if err != nil {
			return nil, err
		}
		encryptedBlocks = append(encryptedBlocks, encryptedBlock)
		message = message[len(block):]
	}
	return encryptedBlocks, nil
}

// EncryptRSA encrypts a message using RSA-OAEP
func EncryptRSA(publicKey *rsa.PublicKey, message []byte) (string, error) {
	label := []byte("")  // Optional label for OAEP encoding, can be empty
	hash := sha256.New() // Hash function for OAEP encoding, SHA-256 is commonly used
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, publicKey, message, label)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}
